import argparse
import os
import threading
import time
from pathlib import Path
import cv2
import torch
import numpy as np
import torch.backends.cudnn as cudnn
from playsound import playsound
from threading import Thread
from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
from utils.plots import colors, plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized

without_mask = './load_file/music/without_mask.mp3'
mask_weared_incorrect = './load_file/music/mask_weared_incorrect.mp3'
with_mask = './load_file/music/with_mask.mp3'

def Play(flag=-1, nc=3):
    if nc == 3:
        if flag == 0:
            playsound(without_mask)
        elif flag == 1:
            playsound(mask_weared_incorrect)
        elif flag == 2:
            playsound(with_mask)
    elif nc == 2:
        if flag == 0:
            playsound(without_mask)
        elif flag == 1:
            playsound(with_mask)

@torch.no_grad()    # 以下函数不需要计算梯度
def detect(opt):
    global play_label_flag
    play_label_flag = -1
    weights, view_img, save_txt, imgsz, play_sound = opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.play_sound
    if type(opt.list) is type(''):
        opt.list = [opt.list]
    for source_ele in opt.list:
        source = source_ele
        save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
        webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
            ('rtsp://', 'rtmp://', 'http://', 'https://'))
        print(source)
        # Directories
        save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)  # increment run
        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

        # Initialize
        set_logging()
        device = select_device(opt.device)
        half = device.type != 'cpu'  # half precision only supported on CUDA

        # Load model
        model = attempt_load(weights, map_location=device)  # load FP32 model
        stride = int(model.stride.max())  # model stride
        imgsz = check_img_size(imgsz, s=stride)  # check img_size
        names = model.module.names if hasattr(model, 'module') else model.names  # get class names
        if half:
            model.half()  # to FP16

        # Second-stage classifier
        classify = False
        if classify:
            modelc = load_classifier(name='resnet101', n=2)  # initialize
            modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

        # Set Dataloader
        vid_path, vid_writer = None, None
        if webcam:
            view_img = check_imshow()
            cudnn.benchmark = True  # set True to speed up constant image size inference
            dataset = LoadStreams(source, img_size=imgsz, stride=stride)
        else:
            dataset = LoadImages(source, img_size=imgsz, stride=stride)

        # Run inference
        if device.type != 'cpu':
            model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
        t0 = time.time()

        file_name = 0
        # 摄像头获取图像循环
        for path, img, im0s, vid_cap in dataset:

            img = torch.from_numpy(img).to(device)

            img = img.half() if half else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0

            if img.ndimension() == 3:
                img = img.unsqueeze(0)
            # Inference
            t1 = time_synchronized()
            pred = model(img, augment=opt.augment)[0]

            # Apply NMS
            pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
                                       max_det=opt.max_det)
            t2 = time_synchronized()

            # Apply Classifier
            if classify:
                pred = apply_classifier(pred, modelc, img, im0s)
            # Process detections
            for i, det in enumerate(pred):  # detections per image
                if webcam:  # batch_size >= 1
                    p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
                else:
                    p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)

                p = Path(p)  # to Path
                save_path = str(save_dir / p.name)  # img.jpg
                txt_path = str(save_dir / 'labels' / p.stem) + (
                    '' if dataset.mode == 'image' else f'_{frame}')  # img.txt
                s += '%gx%g ' % img.shape[2:]  # print string
                gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                imc = im0.copy() if opt.save_crop else im0  # for opt.save_crop
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    # 检测目标数
                    object_nums = 0
                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                        object_nums += 1
                        if save_txt:  # Write to file
                            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                            line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
                            with open(txt_path + '.txt', 'a') as f:
                                f.write(('%g ' * len(line)).rstrip() % line + '\n')

                        if save_img or opt.save_crop or view_img:  # Add bbox to image
                            c = int(cls)  # integer class
                            # label = label + conf
                            label = None if opt.hide_labels else (
                                names[c] if opt.hide_conf else f'{names[c]} {conf:.2f}')
                            # 播放声音
                            if play_sound == 'open':
                                # 仅目标数为1播放声音
                                if object_nums == 1:
                                    # 播放声音(2classes)
                                    # if play_label_flag == -1:
                                    #     if label.split(' ')[0] == 'without_mask':
                                    #         play_label_flag = 0
                                    #         Thread(target=Play, args=(0 ,2 ,)).start()
                                    #     elif label.split(' ')[0] == 'with_mask':
                                    #         play_label_flag = 1
                                    #         Thread(target=Play, args=(1 ,2 ,)).start()
                                    # elif play_label_flag == 0:
                                    #     if label.split(' ')[0] == 'with_mask':
                                    #         play_label_flag = 1
                                    #         Thread(target=Play, args=(1 ,2 ,)).start()
                                    # elif play_label_flag == 1:
                                    #     if label.split(' ')[0] == 'without_mask':
                                    #         play_label_flag = 0
                                    #         Thread(target=Play, args=(0 ,2 ,)).start()

                                    # 播放声音(3classes)
                                    if play_label_flag == -1:
                                        if label.split(' ')[0] == 'without_mask':
                                            play_label_flag = 0
                                            Thread(target=Play, args=(0, 3,)).start()
                                        elif label.split(' ')[0] == 'mask_weared_incorrect':
                                            play_label_flag = 1
                                            # Thread(target=Play, args=(1, 3,)).start()
                                        elif label.split(' ')[0] == 'with_mask':
                                            play_label_flag = 2
                                            # Thread(target=Play, args=(2, 3,)).start()
                                    elif play_label_flag == 0:
                                        if label.split(' ')[0] == 'mask_weared_incorrect':
                                            play_label_flag = 1
                                            # Thread(target=Play, args=(1, 3,)).start()
                                        elif label.split(' ')[0] == 'with_mask':
                                            play_label_flag = 2
                                            # Thread(target=Play, args=(2, 3,)).start()
                                    elif play_label_flag == 1:
                                        if label.split(' ')[0] == 'without_mask':
                                            play_label_flag = 0
                                            Thread(target=Play, args=(0, 3,)).start()
                                        elif label.split(' ')[0] == 'with_mask':
                                            play_label_flag = 2
                                            # Thread(target=Play, args=(2, 3,)).start()
                                    elif play_label_flag == 2:
                                        if label.split(' ')[0] == 'without_mask':
                                            play_label_flag = 0
                                            Thread(target=Play, args=(0, 3,)).start()
                                        elif label.split(' ')[0] == 'mask_weared_incorrect':
                                            play_label_flag = 1
                                            # Thread(target=Play, args=(1, 3,)).start()

                            plot_one_box(xyxy, im0, label=label, color=colors(c, True),
                                         line_thickness=opt.line_thickness)
                            if opt.save_crop:
                                save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
                else:
                    play_label_flag = -1
                # Print time (inference + NMS)
                print(f'{s}Done. ({t2 - t1:.3f}s)')
                # Stream results
                if view_img:
                    file_path = './temp_images/' + str(file_name).zfill(8) + '.jpg'
                    if not os.path.exists('./temp_images/'):
                        os.mkdir('./temp_images/')
                    cv2.imwrite(file_path, im0)
                    file_name += 1
                    try:
                        # 为节省空间删除前面保存的照片
                        if file_name >= 8:
                            if os.path.exists('./temp_images/' + str(file_name - 8).zfill(8) + '.jpg'):
                                os.remove('./temp_images/' + str(file_name - 8).zfill(8) + '.jpg')
                    except (Exception, UnicodeDecodeError):
                        pass

                    # 480, 640, 3
                    if opt.display_result == 'open':
                        cv2.imshow(str(p), im0)
                        cv2.waitKey(1)  # 1 millisecond

                # Save results (image with detections)
                if save_img:
                    if dataset.mode == 'image':
                        cv2.imwrite(save_path, im0)
                    else:  # 'video' or 'stream'
                        if vid_path != save_path:  # new video
                            vid_path = save_path
                            if isinstance(vid_writer, cv2.VideoWriter):
                                vid_writer.release()  # release previous video writer
                            if vid_cap:  # video
                                fps = vid_cap.get(cv2.CAP_PROP_FPS)
                                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                            else:  # stream
                                fps, w, h = 30, im0.shape[1], im0.shape[0]
                                save_path += '.mp4'
                            vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                        vid_writer.write(im0)

        if save_txt or save_img:
            s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
            print(f"Results saved to {save_dir}{s}")

        print(f'Done. ({time.time() - t0:.3f}s)')


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', '-w', nargs='+', type=str, default='runs/train/exp9/weights/best.pt',
                        help='model.pt path(s)')
    parser.add_argument('--source', '-s', dest="list", nargs='+', type=str, default='0',
                        help='source')  # file/folder, 0 for webcam
    parser.add_argument('--display-result', '-dr', type=str, default='close',
                        help='display detect result video, open 打开, close 关闭')
    parser.add_argument('--play-sound', '-ps', type=str, default='close',
                        help='play with_mask or without_mask sound etc., open 打开, close 关闭')
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
    parser.add_argument('--device', '-d', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    opt = parser.parse_args()
    print(opt)
    check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))

    if opt.update:  # update all models (to fix SourceChangeWarning)
        for opt.weights in ['yolov3.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt']:
            detect(opt=opt)
            strip_optimizer(opt.weights)
    else:
        detect(opt=opt)
